package org.acghub.langchain4j.controller;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import lombok.extern.slf4j.Slf4j;
import org.acghub.langchain4j.message.DefaultMessageMemory;
import org.springframework.util.StopWatch;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.concurrent.TimeUnit;

@RestController
@Slf4j
public class ChatController {

    private final DefaultMessageMemory<String, ChatMessage> langChainChatMessageMemory = new DefaultMessageMemory<>(16);
    ChatModel chatModel;

    public ChatController(ChatModel chatModel) {
        this.chatModel = chatModel;
    }

    @GetMapping("/chat")
    public String chat(@RequestParam(value = "uid", defaultValue = "robot") String uid, @RequestParam(value = "msg", defaultValue = "Hello") String msg) {
        StopWatch stopWatch = new StopWatch();
        stopWatch.start();

        UserMessage userMessage = UserMessage.from(msg);

        // 获取历史请求
        List<ChatMessage> history = langChainChatMessageMemory.getHistory(uid);
        history.add(userMessage);

        ChatResponse chatResponse = chatModel.chat(history);
        AiMessage aiMessage = chatResponse.aiMessage();
        String resp = aiMessage.text();


        stopWatch.stop();
        stopWatch.prettyPrint(TimeUnit.NANOSECONDS);

        // 保存历史请求
        langChainChatMessageMemory.addLast(uid, userMessage);
        langChainChatMessageMemory.addLast(uid, aiMessage);

        log.info("[chat]ChatController.chat@|msg={} | resp={}", msg, resp);
        return resp;
    }
}